import tool as tl
import os
import numpy as np
labelbar = {
    'baseline 2layer cnn': 'baseline',
    '4layer_cnn_0.7': '4-layer KD',
    '6layer_cnn_0.78': '6-layer KD',
    '8layer_cnn_0.81': '8-layer KD',
    '10layer_cnn_0.84': '10-layer KD',
    'lsr kd cnn': 'LSR KD',
    'fake noise teacher': 'FNT KD'
}
def loss_picnum(path):
    class_names = os.listdir(path)

    curvedatas = []
    label = []
    for classname in class_names:
        # if not('baseline' in classname):
        #     continue
        if (not 'fake' in classname) and (not '6' in classname):
            continue
        label.append(labelbar[classname])
        print('{} is reading...'.format(classname))
        classpath = path + classname + '/'
        losspath = classpath + 'loss/'
        loss_lists = os.listdir(losspath)

        lossmatrix = []

        for i in range(len(loss_lists)):
            loss = (tl.to_float_dim1(tl.read_csv(losspath + loss_lists[i])))
            lossmatrix.append(loss)

        ave_loss = tl.get_mean_arr(np.array(lossmatrix))


        curvedatas.append(ave_loss)
        print(ave_loss[-1])

        print('{} read finish'.format(classname))
    return curvedatas, label

#
# path = '../weights_params/IE demo val loss 160 epoch/'
#
# tl.pic_make_matrix(curvedatas, label, '', 1)
